# 从ultralytics库中导入YOLO模型
from ultralytics import YOLO  
# 从datasets库中导入load_dataset方法用于加载数据集
from datasets import load_dataset   
# 从PIL图像处理库中导入Image类
from PIL import Image   

# 使用load_dataset方法加载塑料垃圾检测数据集
dataset =load_dataset("kili-technology/plastic_in_river")
# 从数据集中取出一张测试图片    
img = dataset["test"][0]["image"]
# 加载已训练好的YOLO检测模型
model =YOLO("runs/detect/train/weights/best.pt") 
# 使用模型对图片进行预测,结果保存在res变量中  
res = model.predict(img)[0]
# 打印出预测框坐标
print(res.boxes)
# 可视化预测结果
res = res.plot(line_width=1)  
# 将RGB通道顺序翻转成BGR顺序以用于opencv显示
res = res[:, :, ::-1]
# 将预测的可视化图像转换为PIL Image对象  
res = Image.fromarray(res)
# 保存预测visualization图像
res.save("output.png")